from python_ai.common.xcommon import sep
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import sys
import os
import datetime

sep('Check version')
print(tf.__version__)

sep('Load data')
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
for d in (train_images, train_labels, test_images, test_labels):
    print(d.shape, d.dtype)
print(np.unique(train_labels))
print(np.unique(test_labels))
# print(train_images[0])
class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
class_idx = range(len(class_names))
for i, cls_name in enumerate(class_names):
    print(f'{i}: {cls_name}', end=', ')
print()

sep('Visualize data')
plt.figure(figsize=[12, 6])
spr = 2
spc = 4
spn = 0
for i in range(spr*spc):
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.imshow(train_images[i])
    plt.colorbar()
    plt.grid(False)
plt.show(block=False)

sep('Preprocess data')
# train_images /= 255.0  # ValueError: output array is read-only
train_images = train_images / 255.0
test_images = test_images / 255.0
for d in (train_images, train_labels, test_images, test_labels):
    print(d.shape, d.dtype)

sep('Visualize data')
plt.figure(figsize=[12, 12])
spr = 5
spc = 5
spn = 0
for i in range(spr*spc):
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    plt.imshow(train_images[i], cmap=plt.cm.binary)
    plt.xlabel(class_names[train_labels[i]])
plt.show(block=False)

sep('Model')
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
epochs = 2
ver = 'v1.0_epochs'+ str(epochs)
path = sys.argv[0] + '_' + ver + '.tmp.dat'
if os.path.exists(path):
    sep('Loading saved model weight')
    model.load_weights(path)
else:
    sep('Train the model')
    dt1 = datetime.datetime.now()
    model.fit(train_images, train_labels, epochs=epochs)
    dt2 = datetime.datetime.now()
    dt_diff = dt2 - dt1
    print(f'Time diff: {dt_diff}')
    sep('Save the trained model weights')
    model.save_weights(path, save_format='h5')

sep('Evaluate accuracy')
test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)
print(f'test_loss={test_loss}, test_acc={test_acc}')

sep('Make predictions')
proba_model = tf.keras.Sequential([
    model,
    tf.keras.layers.Softmax()
])
predictions = proba_model.predict(test_images)
plt.figure(figsize=[12, 12])
spr = 2
spc = 2
spn = 0
for i in range(spr*spc):
    spn += 1
    plt.subplot(spr, spc, spn)
    pred = predictions[i]
    plt.plot(class_idx, pred)
    plt.xlabel(class_names[test_labels[i]])

sep('End')
# Finally show all plotting
plt.show()
